package online.heycm.platform.canary.gateway;

import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

import lombok.extern.slf4j.Slf4j;
import online.heycm.platform.canary.common.CanaryConstant;
import online.heycm.platform.canary.common.CanaryHolder;
import online.heycm.platform.canary.common.CanaryLoadBalancer;
import online.heycm.platform.canary.common.CanaryRule;
import online.heycm.platform.canary.common.CanaryRuleCache;
import online.heycm.platform.canary.common.CanaryStrategy;
import org.springframework.beans.factory.ObjectProvider;
import org.springframework.cloud.client.ServiceInstance;
import org.springframework.cloud.client.loadbalancer.reactive.Request;
import org.springframework.cloud.loadbalancer.core.ServiceInstanceListSupplier;
import org.springframework.http.HttpHeaders;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.util.CollectionUtils;
import org.springframework.util.StringUtils;
import org.springframework.web.server.ServerWebExchange;

/**
 * 网关灰度负载均衡实现
 *
 * @author heycm
 * @since 2023/12/30 9:11
 */
@Slf4j
public class GatewayCanaryLoadBalancer extends CanaryLoadBalancer {

    public GatewayCanaryLoadBalancer(ObjectProvider<ServiceInstanceListSupplier> serviceInstanceListSupplierProvider,
                                     String serviceId) {
        super(serviceInstanceListSupplierProvider, serviceId);
        log.info("GatewayCanaryLoadBalancer init...");
    }

    public GatewayCanaryLoadBalancer(ObjectProvider<ServiceInstanceListSupplier> serviceInstanceListSupplierProvider,
                                     String serviceId, int seedPosition) {
        super(serviceInstanceListSupplierProvider, serviceId, seedPosition);
        log.info("GatewayCanaryLoadBalancer init...");
    }

    private List<ServiceInstance> getNormal(List<ServiceInstance> instances, String path) {
        List<ServiceInstance> normalInstances = CanaryHolder.getAllNormalInstances(instances);
        if (CollectionUtils.isEmpty(normalInstances)) {
            log.warn("请求 {} 无正式服务!!!", path);
            return null;
        }
        return normalInstances;
    }

    @Override
    protected List<ServiceInstance> canaryChoose(Request request, List<ServiceInstance> instances) {
        // 获取请求路由
        ServerWebExchange exchange = (ServerWebExchange) request.getContext();
        ServerHttpRequest serverHttpRequest = exchange.getRequest();
        String path = serverHttpRequest.getURI().getPath();

        // 未开启灰度配置，优先转发正常服务
        if (!CanaryRuleCache.enable()) {
            return getNormal(instances, path);
        }

        // 匹配灰度配置
        List<CanaryRule> rules = CanaryHolder.match(getServiceId(), path);

        // 不匹配灰度时，优先转发正常服务
        if (CollectionUtils.isEmpty(rules)) {
            return getNormal(instances, path);
        }

        // 匹配灰度规则
        CanaryRule rule = CanaryHolder.highestPriority(rules);
        CanaryStrategy strategy = CanaryStrategy.ofNullable(rule.getCanaryStrategy());

        // 特定请求头灰度
        if (strategy == CanaryStrategy.HEAD) {
            return headGrayStrategy(path, serverHttpRequest, rule, instances);
        }
        // 流量灰度
        if (strategy == CanaryStrategy.FLOW) {
            return flowGrayStrategy(path, exchange, rule, instances);
        }

        return null;
    }

    private List<ServiceInstance> headGrayStrategy(String path, ServerHttpRequest request, CanaryRule rule, List<ServiceInstance> instances) {
        // 取请求头参数值 X-CANARY-HEAD，取不到请求头时，优先转发正常服务
        HttpHeaders headers = request.getHeaders();
        if (CollectionUtils.isEmpty(headers)) {
            return getNormal(instances, path);
        }
        List<String> canaryHead = headers.get(CanaryStrategy.HEAD.getHeader());
        if (CollectionUtils.isEmpty(canaryHead) || StringUtils.isEmpty(canaryHead.get(0))) {
            return getNormal(instances, path);
        }
        String canaryValue = canaryHead.get(0);

        // 请求头的值与灰度配置值对比，不匹配时，优先转发正常服务
        if (StringUtils.isEmpty(canaryValue) || !rule.getCanaryValue().contains(canaryValue)) {
            log.info("请求 {}:{} {} 不匹配灰度配置，不做灰度分发", CanaryStrategy.HEAD.getHeader(), canaryValue, path);
            return getNormal(instances, path);
        }

        // 过滤匹配的灰度实例
        List<ServiceInstance> grayInstances = CanaryHolder.getAllCanaryInstances(instances).stream().filter(inst -> {
            Map<String, String> metadata = inst.getMetadata();
            if (CollectionUtils.isEmpty(metadata)) {
                return false;
            }
            String metadataKey = metadata.get(CanaryStrategy.HEAD.getMetadata());
            return metadataKey != null && metadataKey.contains(canaryValue);
        }).collect(Collectors.toList());

        // 没有灰度实例时，优先转发正常服务
        if (CollectionUtils.isEmpty(grayInstances)) {
            log.info("请求 {}:{} {} 无匹配灰度服务，不做灰度分发", CanaryStrategy.HEAD.getHeader(), canaryValue, path);
            return getNormal(instances, path);
        }
        return grayInstances;
    }

    private List<ServiceInstance> flowGrayStrategy(String path, ServerWebExchange exchange, CanaryRule rule, List<ServiceInstance> instances) {
        // 过滤匹配的灰度实例
        List<ServiceInstance> grayInstances = CanaryHolder.getAllCanaryInstances(instances).stream().filter(inst -> {
            Map<String, String> metadata = inst.getMetadata();
            if (CollectionUtils.isEmpty(metadata)) {
                return false;
            }
            String metadataKey = metadata.get(CanaryStrategy.FLOW.getMetadata());
            return CanaryConstant.METADATA_FLOW_VALUE.equals(metadataKey);
        }).collect(Collectors.toList());

        // 是否走流量灰度
        boolean markFlow = CanaryHolder.markFlow(Integer.parseInt(rule.getCanaryValue()));

        // 不走流量灰度时，优先转发正常服务
        if (!markFlow) {
            return getNormal(instances, path);
        }

        // 设置请求头参数值 X-CANARY-FLOW
        ServerHttpRequest serverHttpRequest = exchange.getRequest().mutate()
                .header(CanaryStrategy.FLOW.getHeader(), CanaryConstant.CANARY_FLOW_VALUE)
                .build();
        exchange = exchange.mutate().request(serverHttpRequest).build();

        // 没有灰度实例时，转发到普通实例
        if (CollectionUtils.isEmpty(grayInstances)) {
            log.info("请求 {}:{} {} 无匹配灰度服务，不做灰度分发", CanaryStrategy.FLOW.getHeader(), CanaryConstant.CANARY_FLOW_VALUE, path);
            return instances;
        }
        return grayInstances;
    }
}
